# functions that implement analysis and synthesis of sounds using the Stochastic Model
# (for example usage check stochasticModel_function.py in the interface directory)

import numpy as np
from scipy.fft import fft, ifft
from scipy.interpolate import splev, splrep
from scipy.signal import resample
from scipy.signal.windows import hann

from smstools.models import utilFunctions as UF


def hertz_to_mel(f):
    """
    Conversion from hertz scale to mel scale
    """
    return 2595 * np.log10(1 + f / 700)


def mel_to_hetz(m):
    """
    Conversion from mel scale to hertz scale
    """
    return 700 * (10 ** (m / 2595) - 1)


def stochasticModelAnal(x, H, N, stocf, fs=44100, melScale=1):
    """
    Stochastic analysis of a sound
    x: input array sound, H: hop size, N: fftsize
    stocf: decimation factor of mag spectrum for stochastic analysis, bigger than 0, maximum of 1
    fs: sampling rate
    melScale: choose between linear scale, 0, or mel scale, 1
    returns stocEnv: stochastic envelope
    """

    hN = N // 2 + 1  # positive size of fft
    No2 = N // 2  # half of N
    if hN * stocf < 3:  # raise exception if decimation factor too small
        raise ValueError("Stochastic decimation factor too small")

    if stocf > 1:  # raise exception if decimation factor too big
        raise ValueError("Stochastic decimation factor above 1")

    if H <= 0:  # raise error if hop size 0 or negative
        raise ValueError("Hop size (H) smaller or equal to 0")

    if not (UF.isPower2(N)):  # raise error if N not a power of two
        raise ValueError("FFT size (N) is not a power of 2")

    w = hann(N)  # analysis window
    x = np.append(
        np.zeros(No2), x
    )  # add zeros at beginning to center first window at sample 0
    x = np.append(x, np.zeros(No2))  # add zeros at the end to analyze last sample
    pin = No2  # initialize sound pointer in middle of analysis window
    pend = x.size - No2  # last sample to start a frame
    if melScale == 1:
        binFreqsMel = hertz_to_mel(np.arange(hN) * fs / float(N))
        uniformMelFreq = np.linspace(binFreqsMel[0], binFreqsMel[-1], hN)
    while pin <= pend:
        xw = x[pin - No2 : pin + No2] * w  # window the input sound
        X = fft(xw)  # compute FFT
        mX = 20 * np.log10(abs(X[:hN]))  # magnitude spectrum of positive frequencies
        if melScale == 1:
            spl = splrep(binFreqsMel, np.maximum(-200, mX))
            mY = resample(
                splev(uniformMelFreq, spl), int(stocf * hN)
            )  # decimate the mag spectrum
        else:
            mY = resample(
                np.maximum(-200, mX), int(stocf * hN)
            )  # decimate the mag spectrum
        if pin == No2:  # first frame
            stocEnv = np.array([mY])
        else:  # rest of frames
            stocEnv = np.vstack((stocEnv, np.array([mY])))
        pin += H  # advance sound pointer
    return stocEnv


def stochasticModelSynth(stocEnv, H, N, fs=44100, melScale=1):
    """
    Stochastic synthesis of a sound
    stocEnv: stochastic envelope; H: hop size; N: fft size
    fs: sampling rate
    melScale: choose between linear scale, 0, or mel scale, 1 (should match the analysis)
    returns y: output sound
    """

    if not (UF.isPower2(N)):  # raise error if N not a power of two
        raise ValueError("N is not a power of two")

    hN = N // 2 + 1  # positive size of fft
    No2 = N // 2  # half of N
    L = stocEnv[:, 0].size  # number of frames
    ysize = H * (L + 3)  # output sound size
    y = np.zeros(ysize)  # initialize output array
    ws = 2 * hann(N)  # synthesis window
    pout = 0  # output sound pointer
    if melScale == 1:
        binFreqsMel = hertz_to_mel(np.arange(hN) * fs / float(N))
        uniformMelFreq = np.linspace(binFreqsMel[0], binFreqsMel[-1], hN)
    for l in range(L):
        mY = resample(stocEnv[l, :], hN)  # interpolate to original size
        if melScale == 1:
            spl = splrep(uniformMelFreq, mY)
            mY = splev(binFreqsMel, spl)
        pY = 2 * np.pi * np.random.rand(hN)  # generate phase random values
        Y = np.zeros(N, dtype=complex)  # initialize synthesis spectrum
        Y[:hN] = 10 ** (mY / 20) * np.exp(1j * pY)  # generate positive freq.
        Y[hN:] = 10 ** (mY[-2:0:-1] / 20) * np.exp(
            -1j * pY[-2:0:-1]
        )  # generate negative freq.
        fftbuffer = np.real(ifft(Y))  # inverse FFT
        y[pout : pout + N] += ws * fftbuffer  # overlap-add
        pout += H
    y = np.delete(y, range(No2))  # delete half of first window
    y = np.delete(y, range(y.size - No2, y.size))  # delete half of the last window
    return y


def stochasticModel(x, H, N, stocf, fs=44100, melScale=1):
    """
    Stochastic analysis/synthesis of a sound, one frame at a time
    x: input array sound, H: hop size, N: fft size
    stocf: decimation factor of mag spectrum for stochastic analysis, bigger than 0, maximum of 1
    fs: sampling rate
    melScale: choose between linear scale, 0, or mel scale, 1 (should match the analysis)
    returns y: output sound
    """
    hN = N // 2 + 1  # positive size of fft
    No2 = N // 2  # half of N
    if hN * stocf < 3:  # raise exception if decimation factor too small
        raise ValueError("Stochastic decimation factor too small")

    if stocf > 1:  # raise exception if decimation factor too big
        raise ValueError("Stochastic decimation factor above 1")

    if H <= 0:  # raise error if hop size 0 or negative
        raise ValueError("Hop size (H) smaller or equal to 0")

    if not (UF.isPower2(N)):  # raise error if N not a power of twou
        raise ValueError("FFT size (N) is not a power of 2")

    w = hann(N)  # analysis/synthesis window
    x = np.append(
        np.zeros(No2), x
    )  # add zeros at beginning to center first window at sample 0
    x = np.append(x, np.zeros(No2))  # add zeros at the end to analyze last sample
    pin = No2  # initialize sound pointer in middle of analysis window
    pend = x.size - No2  # last sample to start a frame
    y = np.zeros(x.size)  # initialize output array
    if melScale == 1:
        binFreqsMel = hertz_to_mel(np.arange(hN) * fs / float(N))
        uniformMelFreq = np.linspace(binFreqsMel[0], binFreqsMel[-1], hN)
    while pin <= pend:
        # -----analysis-----
        xw = x[pin - No2 : pin + No2] * w  # window the input sound
        X = fft(xw)  # compute FFT
        mX = 20 * np.log10(abs(X[:hN]))  # magnitude spectrum of positive frequencies
        if melScale == 1:
            spl = splrep(binFreqsMel, np.maximum(-200, mX))
            stocEnv = resample(
                splev(uniformMelFreq, spl), int(stocf * hN)
            )  # decimate the mag spectrum
        else:
            stocEnv = resample(
                np.maximum(-200, mX), int(stocf * hN)
            )  # decimate the mag spectrum
        # -----synthesis-----
        mY = resample(stocEnv, hN)  # interpolate to original size
        if melScale == 1:
            spl = splrep(uniformMelFreq, mY)
            mY = splev(binFreqsMel, spl)
        pY = 2 * np.pi * np.random.rand(hN)  # generate phase random values
        Y = np.zeros(N, dtype=complex)
        Y[:hN] = 10 ** (mY / 20) * np.exp(1j * pY)  # generate positive freq.
        Y[hN:] = 10 ** (mY[-2:0:-1] / 20) * np.exp(
            -1j * pY[-2:0:-1]
        )  # generate negative freq.
        fftbuffer = np.real(ifft(Y))  # inverse FFT
        y[pin - No2 : pin + No2] += w * fftbuffer  # overlap-add
        pin += H  # advance sound pointer
    y = np.delete(y, range(No2))  # delete half of first window which was added
    y = np.delete(
        y, range(y.size - No2, y.size)
    )  # delete half of last window which was added
    return y
